package com.lw.leetcode.linked.b;

import com.lw.leetcode.tree.TreeNode;

import java.util.*;

/**
 * 1110. 删点成林
 *
 * @Author liw
 * @Date 2021/10/7 19:58
 * @Version 1.0
 */
public class DelNodes {

    public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
        Queue<TreeNode> queue = new LinkedList<>();
        List<TreeNode> list = new ArrayList<>();
        Set<Integer> set = new HashSet<>(to_delete.length);
        for (int i : to_delete) {
            set.add(i);
        }
        queue.add(root);
        if (!set.contains(root.val)) {
            list.add(root);
        }
        while (!queue.isEmpty()) {
            TreeNode p = queue.poll();
            if (p.left != null) {
                queue.add(p.left);
                if (set.contains(p.val) && !set.contains(p.left.val)) {
                    list.add(p.left);
                }
                if (!set.contains(p.val) && set.contains(p.left.val)) {
                    p.left = null;
                }
            }
            if (p.right != null) {
                queue.add(p.right);
                if (set.contains(p.val) && !set.contains(p.right.val)) {
                    list.add(p.right);
                }
                if (!set.contains(p.val) && set.contains(p.right.val)) {
                    p.right = null;
                }
            }
        }
        return list;
    }

}
